"""
@Time: 2020/12/4 上午 11:17
@Author: jinzhuan
@File: ner_predictor_test.py
@Desc: 
"""
import torch
from cognlp import *
from cognlp.io.processor.ner.conll2003 import TrexNerProcessor


torch.cuda.set_device(4)
device = torch.device('cuda')
processor = TrexNerProcessor(path='../data/ner/trex/data')
model = Bert4Ner(len(processor.vocabulary), device=device)
predictor = NerPredictor(model, model_path='../data/ner/trex/model/2020-12-04-03:46:31-model.pkl',
                         vocabulary=processor.vocabulary, device=device, device_ids=[4, 5, 6])
print(predictor.predict("located in the eastern Pyrenees mountains and bordered by Spain and France. Created under a charter in A.D. 988, the present Principality was formed in A.D. 1278. It is known as a principality as it is a monarchy headed by two Co-Princes – the Spanish/Roman Catholic Bishop of Urgell and the President of France.")[1])
